import random
import numpy as np

import paddle.fluid as fluid
from paddle.fluid.dygraph import ParallelEnv
from helm.static.models.layers import DEFAULTS

__all__ = ["set_device", "set_seed"]


def set_device(device):
    """
    Args:
        device (str): specify device type, 'cpu' or 'gpu'.

    Returns:
        fluid.CUDAPlace or fluid.CPUPlace: Created GPU or CPU place.
    """

    assert isinstance(device, str) and device.lower() in ['cpu', 'gpu'], \
        "Expected device in ['cpu', 'gpu'], but got {}".format(device)

    place = fluid.CUDAPlace(ParallelEnv().dev_id) \
        if device.lower() == 'gpu' and fluid.is_compiled_with_cuda() \
        else fluid.CPUPlace()

    return place


def set_seed(seed=0):
    DEFAULTS['seed'] = seed
    random.seed(seed)
    np.random.seed(seed)
    fluid.default_main_program().random_seed = seed
    fluid.default_startup_program().random_seed = seed